"""
环境检查脚本
检查项目依赖和数据是否准备就绪
"""
import os
import sys


def check_python_version():
    """检查Python版本"""
    print("检查Python版本...")
    version = sys.version_info
    if version.major >= 3 and version.minor >= 8:
        print(f"  ✓ Python版本: {version.major}.{version.minor}.{version.micro}")
        return True
    else:
        print(f"  ✗ Python版本过低: {version.major}.{version.minor}.{version.micro}")
        print("    需要Python 3.8或更高版本")
        return False


def check_packages():
    """检查必要的包"""
    print("\n检查必要的包...")
    required_packages = [
        'torch',
        'torchvision',
        'numpy',
        'cv2',
        'PIL',
        'sklearn',
        'albumentations',
        'tqdm'
    ]

    all_installed = True

    for package in required_packages:
        try:
            if package == 'cv2':
                __import__('cv2')
            elif package == 'PIL':
                __import__('PIL')
            elif package == 'sklearn':
                __import__('sklearn')
            else:
                __import__(package)
            print(f"  ✓ {package}")
        except ImportError:
            print(f"  ✗ {package} 未安装")
            all_installed = False

    if not all_installed:
        print("\n请运行以下命令安装依赖：")
        print("  pip install -r requirements.txt")

    return all_installed


def check_cuda():
    """检查CUDA是否可用"""
    print("\n检查CUDA...")
    try:
        import torch
        if torch.cuda.is_available():
            print(f"  ✓ CUDA可用")
            print(f"    设备数量: {torch.cuda.device_count()}")
            print(f"    当前设备: {torch.cuda.get_device_name(0)}")
            return True
        else:
            print("  ⚠ CUDA不可用，将使用CPU训练（速度较慢）")
            return False
    except ImportError:
        print("  ✗ PyTorch未安装")
        return False


def check_directories():
    """检查目录结构"""
    print("\n检查目录结构...")
    required_dirs = [
        'baselines',
        'datasets',
        'metrics',
        'utils',
        'configs',
        'data',
        'data/images',
        'data/masks'
    ]

    all_exist = True

    for dir_path in required_dirs:
        if os.path.exists(dir_path):
            print(f"  ✓ {dir_path}/")
        else:
            print(f"  ✗ {dir_path}/ 不存在")
            all_exist = False

    if not all_exist:
        print("\n某些目录不存在，将自动创建...")
        for dir_path in required_dirs:
            os.makedirs(dir_path, exist_ok=True)
        print("  ✓ 目录创建完成")

    return True


def check_data():
    """检查数据"""
    print("\n检查数据...")

    images_dir = 'data/images'
    masks_dir = 'data/masks'

    if not os.path.exists(images_dir):
        print(f"  ✗ {images_dir} 不存在")
        return False

    if not os.path.exists(masks_dir):
        print(f"  ✗ {masks_dir} 不存在")
        return False

    # 统计图像数量
    image_files = [f for f in os.listdir(images_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]
    mask_files = [f for f in os.listdir(masks_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]

    print(f"  图像数量: {len(image_files)}")
    print(f"  标签数量: {len(mask_files)}")

    if len(image_files) == 0:
        print("  ⚠ 没有找到图像文件")
        print("    请将图像放入 data/images/ 目录")
        return False

    if len(mask_files) == 0:
        print("  ⚠ 没有找到标签文件")
        print("    请将标签放入 data/masks/ 目录")
        return False

    if len(image_files) != len(mask_files):
        print("  ⚠ 图像和标签数量不匹配")

    # 检查文件名是否对应
    image_names = set([os.path.splitext(f)[0] for f in image_files])
    mask_names = set([os.path.splitext(f)[0] for f in mask_files])

    missing_masks = image_names - mask_names
    missing_images = mask_names - image_names

    if missing_masks:
        print(f"  ⚠ {len(missing_masks)} 个图像缺少对应的标签")

    if missing_images:
        print(f"  ⚠ {len(missing_images)} 个标签缺少对应的图像")

    if not missing_masks and not missing_images:
        print("  ✓ 所有图像和标签都有对应")

    return len(image_files) > 0 and len(mask_files) > 0


def check_models():
    """检查模型文件"""
    print("\n检查模型文件...")

    models = ['unet', 'deeplabv3plus', 'fcn', 'pspnet', 'unetpp', 'segnet']

    all_exist = True

    for model in models:
        model_file = os.path.join('baselines', model, 'model.py')
        config_file = os.path.join('baselines', model, 'config.py')
        train_file = os.path.join('baselines', model, 'train.py')
        infer_file = os.path.join('baselines', model, 'infer.py')

        if all([os.path.exists(f) for f in [model_file, config_file, train_file, infer_file]]):
            print(f"  ✓ {model}")
        else:
            print(f"  ✗ {model} 文件不完整")
            all_exist = False

    return all_exist


def main():
    """主函数"""
    print("=" * 80)
    print("雷达俯视图点目标实例分割项目 - 环境检查")
    print("=" * 80)

    results = []

    results.append(("Python版本", check_python_version()))
    results.append(("依赖包", check_packages()))
    results.append(("CUDA", check_cuda()))
    results.append(("目录结构", check_directories()))
    results.append(("数据", check_data()))
    results.append(("模型文件", check_models()))

    print("\n" + "=" * 80)
    print("检查结果汇总")
    print("=" * 80)

    for name, result in results:
        status = "✓" if result else "✗"
        print(f"  {status} {name}")

    all_passed = all([r[1] for r in results])

    print("\n" + "=" * 80)

    if all_passed:
        print("✓ 所有检查通过！可以开始训练")
        print("\n快速开始：")
        print("  1. 运行交互式菜单: python quick_start.py")
        print("  2. 训练单个模型: python baselines/unet/train.py --fold 0")
        print("  3. 训练所有模型: python train_all.py")
    else:
        print("✗ 部分检查未通过，请根据上述提示解决问题")

    print("=" * 80 + "\n")


if __name__ == '__main__':
    main()
